import numpy as np
from datetime import datetime

class Network(object):
    def __init__(self, input_dim, hidden_dim, output_dim, learning_rate):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim

        # 初始化权重矩阵
        # wih: 输入层到隐藏层的权重 (hidden_dim x input_dim)
        self.wih = np.random.normal(0.0, pow(self.input_dim, -0.5), (self.hidden_dim, self.input_dim))
        # who: 隐藏层到输出层的权重 (output_dim x hidden_dim)
        self.who = np.random.normal(0.0, pow(self.hidden_dim, -0.5), (self.output_dim, self.hidden_dim))
        
        # 学习率
        self.lr = learning_rate
        
        # 激活函数 (Sigmoid)
        self.activation_function = lambda x: 1 / (1 + np.exp(-x))
        
        pass
    
    def train(self, inputs_list, targets_list):
        # 转换输入为列向量
        inputs = np.array(inputs_list, ndmin=2).T
        targets = np.array(targets_list, ndmin=2).T
        
        # 前向传播
        # 隐藏层
        hidden_inputs = np.dot(self.wih, inputs)
        hidden_outputs = self.activation_function(hidden_inputs)
        
        # 输出层
        final_inputs = np.dot(self.who, hidden_outputs)
        final_outputs = self.activation_function(final_inputs)
        
        # 计算误差
        output_errors = targets - final_outputs
        # 隐藏层误差
        hidden_errors = np.dot(self.who.T, output_errors)
        
        # 反向传播，更新权重
        # 更新隐藏层到输出层的权重
        self.who += self.lr * np.dot((output_errors * final_outputs * (1 - final_outputs)), 
                                     np.transpose(hidden_outputs))
        # 更新输入层到隐藏层的权重
        self.wih += self.lr * np.dot((hidden_errors * hidden_outputs * (1 - hidden_outputs)), 
                                     np.transpose(inputs))
        
        pass
    
    def predict(self, inputs_list):
        # 将输入列表转换为数组
        inputs = np.array(inputs_list, ndmin=2).T
        
        # 前向传播计算输出
        hidden_inputs = np.dot(self.wih, inputs)
        hidden_outputs = self.activation_function(hidden_inputs)
        
        final_inputs = np.dot(self.who, hidden_outputs)
        final_outputs = self.activation_function(final_inputs)
        
        return final_outputs


    def save_weights(self, file_prefix="mnist_network"):
        """
        将训练好的权重矩阵保存到磁盘
        :param file_prefix: 文件名前缀
        """
        np.save(f"{file_prefix}_wih.npy", self.wih)  # 输入层到隐藏层权重
        np.save(f"{file_prefix}_who.npy", self.who)  # 隐藏层到输出层权重
        print(f"Weights saved as {file_prefix}_wih.npy and {file_prefix}_who.npy")

    def load_weights(self, file_prefix="mnist_network"):
        """
        从磁盘加载权重矩阵
        :param file_prefix: 文件名前缀
        """
        self.wih = np.load(f"{file_prefix}_wih.npy")
        self.who = np.load(f"{file_prefix}_who.npy")
        print(f"Weights loaded from {file_prefix}_wih.npy and {file_prefix}_who.npy")

